package cn.com.zetatech.loader.common.ftp;

import com.jcraft.jsch.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
public class SftpClient extends ConnectionClient {

    private Session session;
    private ChannelSftp channelSftp;

    public SftpClient(String serverIP, int serverPort, String userName, String password) throws Exception {
        if (serverIP == null || serverPort < 0 || userName == null || userName.equals("")) {
            throw new Exception();
        }
        this.serverIP = serverIP;
        this.serverPort = serverPort;
        this.userName = userName;
        this.password = password;
    }

    @Override
    public void login(int timeout) throws Exception {
        JSch jsch = new JSch();
        session = jsch.getSession(userName, serverIP, serverPort);
        session.setPassword(this.password);
        session.setTimeout(timeout);
        session.setConfig("StrictHostKeyChecking", "no");
        if (!session.isConnected()) {
            session.connect();
        }
        channelSftp = (ChannelSftp) session.openChannel("sftp");
        if (!channelSftp.isConnected()) {
            channelSftp.connect();
            log.info("SFTP连接成功");
        }
    }

    @Override
    public InputStream get(String filePath) throws Exception {
        return channelSftp.get(filePath);
    }

    @Override
    public boolean copy(String sourcePath, String targetDir, String fileName) throws Exception {
        channelSftp.cd("/");
        String[] dirs = targetDir.split("/");
        for (String dir : dirs) {
            if (StringUtils.isBlank(dir)) {
                continue;
            }
            mkDir(dir);
        }
        ChannelExec exec = null;
        try {
            exec = (ChannelExec) session.openChannel("exec");
            String command = "cp " + sourcePath + " " + targetDir + "/" + fileName;
            exec.setCommand(command);
            log.info("执行命令：{}", command);
            exec.connect();
        } finally {
            if (exec != null) {
                exec.disconnect();
            }
        }

        return true;
    }

    @Override
    public boolean remove(String filePath) throws Exception {
        channelSftp.rm(filePath);
        return true;
    }

    @Override
    public boolean move(String sourcePath, String targetDir, String fileName) throws Exception {
        channelSftp.cd("/");
        String[] dirs = targetDir.split("/");
        for (String dir : dirs) {
            if (StringUtils.isBlank(dir)) {
                continue;
            }
            mkDir(dir);
        }
        channelSftp.rename(sourcePath, targetDir + "/" + fileName);
        return true;
    }

    private void mkDir(String dir) {
        try {
            channelSftp.cd(dir);
        } catch (SftpException e) {
            log.info("创建SFTP文件夹：{}", dir);
            try {
                channelSftp.mkdir(dir);
                channelSftp.cd(dir);
            } catch (SftpException sftpException) {
                log.error(e.getMessage(), e);
            }
        }
    }

    @Override
    public Map<String, Long> listFileSize(String dir, List<String> fileNames) {
        Map<String, Long> fileSizes = new HashMap<>();
        List<ChannelSftp.LsEntry> sftpFiles = new ArrayList<>();
        try {
            channelSftp.ls(dir, entry -> {
                if (entry.getAttrs().isDir()) {
                    return ChannelSftp.LsEntrySelector.CONTINUE;
                }
                if (fileNames.contains(entry.getFilename())) {
                    sftpFiles.add(entry);
                    return ChannelSftp.LsEntrySelector.CONTINUE;
                }
                if (sftpFiles.size() >= fileNames.size()) {
                    return ChannelSftp.LsEntrySelector.BREAK;
                }
                return ChannelSftp.LsEntrySelector.CONTINUE;
            });
            Map<String, Long> sftpFileNameAndSizes = sftpFiles.stream().collect(Collectors.toMap(ChannelSftp.LsEntry::getFilename, lsEntry -> lsEntry.getAttrs().getSize()));
            for (String fileName : fileNames) {
                fileSizes.put(fileName, sftpFileNameAndSizes.getOrDefault(fileName, 0L));
            }
        } catch (SftpException e) {
            log.error("查询文件大小错误:{}", e.getMessage(), e);
        }
        return fileSizes;
    }

    @Override
    public void disconnect() {
        if (channelSftp != null && channelSftp.isConnected()) {
            channelSftp.exit();
        }
        if (session != null && session.isConnected()) {
            session.disconnect();
        }
        log.info("SFTP连接断开");
    }

    @Override
    public boolean isConnected() {
        return channelSftp != null && channelSftp.isConnected();
    }

    @Override
    public void pendingReset() {

    }
}
